#!/bin/bash

# SecInit - Security Tools Initialization Script

INSTALL_DIR="/usr/local/cloudsec"

# Help function
function show_help() {
    cat <<EOL
Usage: $0 [option]
Options:
  --help           Display this help message.
  --basic          Install and configure basic environment (ohmyzsh, python, java, docker, kubectl, k9s, golang).
    --ohmyzsh      Install oh-my-zsh.
    --python       Install Python.
    --golang       Install Golang.
    --java         Install Java.
    --docker       Install Docker.
    --kubectl      Install kubectl.
    --k9s          Install k9s.
  --pentest        Install penetration testing tools.
    --nmap         Install nmap.
    --metasploit   Install Metasploit.
    --masscan      Install masscan.
    --nuclei       Install Nuclei.
    --oneforall    Install OneForAll.
  --cloudsec       Install cloud security tools.
    --trufflehog   Install trufflehog.
    --cf           Install CF.
    --cdk          Install CDK.
    --trivy        Install trivy.
    --kube-hunter  Install kube-hunter.
    --kubescape    Install kubescape.
  --all            Install all available tools.
EOL
    exit 0
}

# Banner
function show_banner() {
    cat <<EOL
  _________           .__    .___ __________                            
 /   _____/ ____    __|__| __| _/ \______   \__ __  ____   ____ ___.__.
 \_____  \ /  _ \  /  ___/|  |/ |  |   /    / |  |  \/    \ /    <   |  |
 /        (  <_> ) \___ \ |  /  |  |  /    /  |  |  /   |  \   |  \___  |
/_______  /\____/ /____  >|__|  |____/_____ \ |____/|___|  /___|  / ____|
        \/            \/                    \/           \/     \/\/     
EOL
}

# Check if a command exists
function command_exists() {
    command -v "$1" >/dev/null 2>&1
}

# Module 1: Install and configure basic environment
function install_basic_environment() {
    while [[ "$#" -gt 0 ]]; do
        case $1 in
            --ohmyzsh)
                sudo yum install -y zsh git
                sh -c "$(wget -O- https://gitee.com/l0n9/ohmyzsh/raw/master/tools/install.sh)"
                chsh -s "$(command -v zsh)"
                
                git clone https://gitee.com/l0n9/fonts.git && cd fonts && ./install.sh && cd .. && rm -rf fonts
                sed -i 's/robbyrussell/agnoster/g' ~/.zshrc

                plugins_dir=~/.oh-my-zsh/custom/plugins
                git clone https://gitee.com/l0n9/zsh-autosuggestions.git "$plugins_dir/zsh-autosuggestions"
                git clone https://gitee.com/l0n9/zsh-syntax-highlighting.git "$plugins_dir/zsh-syntax-highlighting"

                cat <<EOL >> ~/.zshrc
plugins=(
git
wd
zsh-syntax-highlighting
zsh-autosuggestions
)

source \$ZSH/oh-my-zsh.sh
EOL

                source ~/.zshrc
                echo "oh-my-zsh environment is set up successfully!"

                ;;
            --python)
                echo "Installing Python..."
                sudo yum install -y centos-release-scl scl-utils-build scl-utils
                sudo yum remove -y python3
                sudo yum install -y rh-python38
                source scl_source enable rh-python38
                sudo ln -sf /opt/rh/rh-python38/root/usr/bin/python3.8 /usr/bin/python3
                sudo ln -sf /opt/rh/rh-python38/root/usr/bin/pip3.8 /usr/bin/pip3
                ;;
            --golang)
                echo "Installing Golang..."
                sudo yum -y install bison
                GOLANG_VERSION="1.21.4"
                GOLANG_URL="https://golang.google.cn/dl/go$GOLANG_VERSION.linux-amd64.tar.gz"
                GOLANG_TAR="go$GOLANG_VERSION.linux-amd64.tar.gz"

                wget "$GOLANG_URL" -O "$GOLANG_TAR"
                tar -C /usr/local -xzf "$GOLANG_TAR"
                echo 'export PATH=$PATH:/usr/local/go/bin' >> ~/.zshrc
                source ~/.zshrc
                echo 'export GOPROXY=https://goproxy.cn,direct' >> ~/.zshrc
                rm "$GOLANG_TAR"
                ;;
            --java)
                echo "Installing Java 8..."
                sudo yum -y install java-1.8.0-openjdk-devel
                echo 'export JAVA_HOME=/usr/lib/jvm/java-1.8.0-openjdk' >> ~/.zshrc
                echo 'export PATH=$PATH:$JAVA_HOME/bin' >> ~/.zshrc
                source ~/.zshrc
                ;;
            --docker)
                echo "Installing Docker..."
                sudo yum install -y docker
                sudo systemctl start docker
                sudo systemctl enable docker
                ;;
            --kubectl)
                echo "Installing kubectl..."
                sudo tee /etc/yum.repos.d/kubernetes.repo <<EOF
[kubernetes]
name=Kubernetes
baseurl=https://mirrors.aliyun.com/kubernetes/yum/repos/kubernetes-el7-x86_64/
enabled=1
gpgcheck=1
repo_gpgcheck=0
gpgkey=https://mirrors.aliyun.com/kubernetes/yum/doc/yum-key.gpg https://mirrors.aliyun.com/kubernetes/yum/doc/rpm-package-key.gpg
EOF
                sudo yum install -y kubelet kubeadm kubectl --disableexcludes=kubernetes
                systemctl enable kubelet && systemctl start kubelet
                ;;
            --k9s)
                echo "Installing k9s..."
                git clone https://gitee.com/mirrors/K9s.git "$INSTALL_DIR/k9s"
                cd "$INSTALL_DIR/k9s"
                make build
                echo 'alias k9s="/usr/local/cloudsec/k9s/execs/k9s"' >> ~/.zshrc
                ;;
            *)
                echo "Invalid option for basic environment: $1"
                show_help
                ;;
        esac
        shift
    done

    echo "Basic environment has been installed and configured."
}

# Module 2: Install penetration testing tools
function install_pentest_tools() {
    echo "Installing penetration testing tools..."
    while [[ "$#" -gt 0 ]]; do
        case $1 in
            --nmap)
                echo "Installing nmap..."
                sudo yum install -y nmap
                ;;
            --metasploit)
                echo "Installing Metasploit..."
                curl https://raw.githubusercontent.com/rapid7/metasploit-omnibus/master/config/templates/metasploit-framework-wrappers/msfupdate.erb > msfinstal
                chmod +x msfinstal
                ./msfinstall
                ;;
            --masscan)
                echo "Installing masscan..."
                sudo yum install -y epel-release
                sudo yum install -y masscan
                ;;
            --nuclei)
                echo "Installing Nuclei..."
                git clone https://gitee.com/l0n9/nuclei.git /usr/local/offensive/nuclei
                cd /usr/local/offensive/nuclei/cmd/nuclei
                go build
                sudo mv nuclei /usr/local/bin/
                ;;
            --oneforall)
                echo "Installing OneForAll..."
                git clone https://gitee.com/l0n9/OneForAll.git /usr/local/offensive/OneForAll
                cd /usr/local/offensive/OneForAll
                pip3 install -r requirements.txt
                ;;
            *)
                echo "Invalid option for penetration testing tools: $1"
                show_help
                ;;
        esac
        shift
    done
}

# Module 3: Install cloud security tools
function install_cloudsec_tools() {
    echo "Installing cloud security tools..."
    while [[ "$#" -gt 0 ]]; do
        case $1 in
            --trufflehog)
                echo "Installing trufflehog..."
                git clone https://gitee.com/l0n9/trufflehog.git "$INSTALL_DIR/trufflehog"
                cd "$INSTALL_DIR/trufflehog"
                go install
                ;;
            --cf)
                echo "Installing CF..."
                mkdir -p /usr/local/cloudsec
                wget https://gitee.com/l0n9/sectools/raw/master/cloudsec/cf_linux_amd64 -O /usr/local/cloudsec/cf
                chmod +x /usr/local/cloudsec/cf
                echo 'alias cf="/usr/local/cloudsec/cf"' >> ~/.zshrc
                ;;
            --cdk)
                echo "Installing CDK..."
                CDK_URL="https://gitee.com/l0n9/sectools/raw/master/cloudsec/cdk_linux_amd64"
                wget "$CDK_URL" -O "$INSTALL_DIR/cdk"
                chmod +x "$INSTALL_DIR/cdk"
                echo 'alias cdk="/usr/local/cdk"' >> ~/.zshrc
                ;;
            --trivy)
                echo "Installing trivy..."
                TRIVY_VERSION="v0.47.0"
                TRIVY_URL="https://github.com/aquasecurity/trivy/releases/download/$TRIVY_VERSION/trivy_0.47.0_Linux-64bit.rpm"
                sudo rpm -ivh "$TRIVY_URL"
                ;;
            --kube-hunter)
                echo "Installing kube-hunter..."
                pip3 install kube-hunter
                echo 'alias kube-hunter="/opt/rh/rh-python38/root/usr/local/bin/kube-hunter"' >> ~/.zshrc
                ;;
            --kubescape)
                echo "Installing kubescape..."
                curl -s https://raw.githubusercontent.com/kubescape/kubescape/master/install.sh | /bin/bash
                ;;
            *)
                echo "Invalid option for cloud security tools: $1"
                show_help
                ;;
        esac
        shift
    done

    echo "Cloud security tools have been installed."
}

# Module: Install all available tools
function install_all_tools() {
    install_basic_environment --ohmyzsh --python --java --docker --kubectl --k9s --golang
    install_pentest_tools --nmap --metasploit --masscan --nuclei --oneforall
    install_cloudsec_tools --trufflehog --cf --cdk --trivy --kube-hunter --kubescape
}

# Check for options
if [[ "$1" == "--help" ]]; then
    show_banner
    show_help
fi

# Run the corresponding module based on the input argument
show_banner
case "$1" in
    --basic)
        install_basic_environment "${@:2}"
        ;;
    --pentest)
        install_pentest_tools "${@:2}"
        ;;
    --cloudsec)
        install_cloudsec_tools "${@:2}"
        ;;
    --all)
        install_all_tools
        ;;
    *)
        echo "Invalid option. Please specify a valid module or use --help for assistance."
        show_help
        ;;
esac